# 导入必要的库和模块
from fastapi import APIRouter, Request, HTTPException
from fastapi.responses import StreamingResponse
from dotenv import load_dotenv
from app.services.github_service import GitHubService
from app.services.o3_mini_openai_service import OpenAIo3Service
from app.prompts import (
    SYSTEM_FIRST_PROMPT,
    SYSTEM_SECOND_PROMPT,
    NO_README_INSTRUCTIONS,
)
from app.doc_prompts import (
    SYSTEM_DOC_PROMPT,
    ADDITIONAL_DOC_INSTRUCTIONS_PROMPT,
)
from pydantic import BaseModel
from functools import lru_cache
import json
import asyncio

# 从openai_generate.py导入缓存函数，避免代码重复
from app.routers.openai_generate import get_cached_github_data

# OpenAI文档生成模块
# 本模块负责提供基于OpenAI模型生成项目技术文档的API端点
# 使用三阶段流程生成文档：
# 1. 分析项目结构，生成架构解释
# 2. 映射组件到文件和目录
# 3. 基于前两步结果生成详细的Markdown格式技术文档
# 支持用户自定义指令和私有仓库访问，并提供流式响应实时展示生成过程

load_dotenv()

print("========== 初始化 OpenAI 文档生成路由 ==========")

# 创建FastAPI路由器，设置前缀为/openai/doc，标签为openai-doc
router = APIRouter(prefix="/openai/doc", tags=["openai-doc"])

# 初始化OpenAI o3-mini服务实例
o3_service = OpenAIo3Service()
print("已初始化 OpenAIo3Service 服务")

# 定义API请求体模型
class ApiRequest(BaseModel):
    """API请求体模型，定义了生成文档所需的参数
    
    Attributes:
        username (str): GitHub用户名
        repo (str): 仓库名称
        instructions (str): 自定义指令（可选）
        api_key (str | None): 用户提供的OpenAI API密钥（可选）
        github_pat (str | None): GitHub个人访问令牌（可选）
        model (str): 选择的AI模型，默认为openai-o3
        document_type (str): 文档类型，默认为概要设计文档
    """
    username: str
    repo: str
    instructions: str = ""
    api_key: str | None = None
    github_pat: str | None = None
    model: str = "openai-o3"
    document_type: str = "overview-doc"  # 新增字段，默认为概要设计文档


# 定义计算生成成本的路由
@router.post("/cost")
async def get_doc_generation_cost(request: Request, body: ApiRequest):
    """
    计算生成文档的预估成本
    
    根据仓库文件树和README的令牌数量，计算使用OpenAI API生成文档的预估成本。
    
    Args:
        request (Request): FastAPI请求对象
        body (ApiRequest): 请求体，包含GitHub仓库信息
        
    Returns:
        dict: 包含预估成本或错误信息的字典
    """
    print(f"[OpenAI-Doc] 成本计算请求: 用户名={body.username}, 仓库={body.repo}")
    try:
        # 获取缓存的GitHub数据（文件树和README）
        github_data = get_cached_github_data(body.username, body.repo, body.github_pat)
        file_tree = github_data["file_tree"]
        readme = github_data["readme"]

        # 使用OpenAI服务计算文件树和README的令牌数量
        file_tree_tokens = o3_service.count_tokens(file_tree)
        readme_tokens = o3_service.count_tokens(readme)

        # OpenAI o3-mini模型的成本计算
        # 文档生成通常比图表生成需要更多的输出令牌
        input_cost = ((file_tree_tokens * 2 + readme_tokens) + 3000) * 0.0000011
        output_cost = 16000 * 0.0000044  # 文档输出通常比图表更长
        estimated_cost = input_cost + output_cost

        # 格式化为美元字符串，保留两位小数
        cost_string = f"${estimated_cost:.2f} USD"
        print(f"[OpenAI-Doc] 成本计算结果: {cost_string}")
        return {"cost": cost_string}
    except Exception as e:
        print(f"[OpenAI-Doc] 成本计算错误: {str(e)}")
        return {"error": str(e)}


# 定义流式生成文档的路由
@router.post("/stream")
async def generate_doc_stream(request: Request, body: ApiRequest):
    """
    使用OpenAI生成项目文档并以流式方式返回结果

    Args:
        request (Request): FastAPI请求对象
        body (ApiRequest): 包含GitHub仓库信息和自定义指令的请求体

    Returns:
        StreamingResponse: 流式返回生成的文档内容
    """
    print(f"[OpenAI-Doc] 收到文档生成请求")
    print(f"[OpenAI-Doc] 文档类型详细信息：传入值='{body.document_type}'，类型={type(body.document_type).__name__}")
    print(f"[OpenAI-Doc] 生成请求: 用户名={body.username}, 仓库={body.repo}, 文档类型={body.document_type}, 是否有API密钥={bool(body.api_key)}")
    
    # 在处理开始时打印文档类型相关信息
    print(f"[文档类型-DEBUG-1] 接收到的文档类型: '{body.document_type}'")
    print(f"[文档类型-DEBUG-2] 判断有效类型: {body.document_type in ['overview-doc', 'detailed-doc', 'database-doc', 'api-doc', 'deployment-doc']}")
    
    # 验证指令长度
    if len(body.instructions) > 1000:
        print("[OpenAI-Doc] 错误: 指令长度超过限制")
        return {"error": "Instructions exceed maximum length of 1000 characters"}

    try:
        async def event_generator():
            """
            生成事件流的异步生成器函数
            
            在函数内部包含三个阶段:
            1. 生成项目架构解释
            2. 创建组件映射
            3. 生成详细文档
            
            每个阶段的结果都通过事件流实时返回给客户端
            
            Yields:
                str: 格式化为Server-Sent Events的JSON数据
            """
            try:
                # 获取缓存的GitHub数据
                github_data = get_cached_github_data(
                    body.username, body.repo, body.github_pat
                )
                default_branch = github_data["default_branch"]
                file_tree = github_data["file_tree"]
                readme = github_data["readme"]
                
                # 标记是否存在README
                has_readme = len(readme.strip()) > 0
                if not has_readme:
                    print(f"[OpenAI-Doc] 注意: 仓库 {body.username}/{body.repo} 没有README文件，将仅使用文件树进行分析")
                    # 确保README变量不是None，防止后续计算token出错
                    readme = ""
                
                print(f"[OpenAI-Doc] 已获取GitHub数据: 默认分支={default_branch}, 文件树长度={len(file_tree)}, 有README={has_readme}")
                
                # 发送初始状态
                yield f"data: {json.dumps({'status': 'started', 'message': 'Starting documentation generation...'})}\n\n"
                await asyncio.sleep(0.1)
                
                # 令牌数量检查
                combined_content = f"{file_tree}\n{readme}"
                token_count = o3_service.count_tokens(combined_content)
                print(f"[OpenAI-Doc] 令牌数量检查: 总数={token_count}")

                # 如果令牌数量超过限制且未提供自定义API密钥，返回错误
                if 50000 < token_count < 195000 and not body.api_key:
                    yield f"data: {json.dumps({'error': f'File tree and README combined exceeds token limit (50,000). Current size: {token_count} tokens. This GitHub repository is too large for my wallet, but you can continue by providing your own OpenAI API key.'})}\n\n"
                    return
                # 如果令牌数量超过模型最大上下文长度，返回错误
                elif token_count > 195000:
                    yield f"data: {json.dumps({'error': f'Repository is too large (>195k tokens) for analysis. OpenAI o3-mini\'s max context length is 200k tokens. Current size: {token_count} tokens.'})}\n\n"
                    return

                # 准备提示模板
                first_system_prompt = SYSTEM_FIRST_PROMPT

                # 根据文档类型选择适当的系统提示词
                doc_system_prompt = SYSTEM_DOC_PROMPT
                doc_type_name = "overview-doc"
                
                print(f"[文档类型-DEBUG-3] 选择提示词之前，文档类型值: '{body.document_type}'")
                
                if body.document_type == "detailed-doc":
                    doc_system_prompt = SYSTEM_DETAILED_DOC_PROMPT
                    doc_type_name = "detailed-doc"
                    print(f"[OpenAI-Doc] 文档类型: 详细设计文档 (detailed-doc)")
                elif body.document_type == "database-doc":
                    doc_system_prompt = SYSTEM_DATABASE_DOC_PROMPT
                    doc_type_name = "database-doc"
                    print(f"[OpenAI-Doc] 文档类型: 数据库文档 (database-doc)")
                elif body.document_type == "api-doc":
                    doc_system_prompt = SYSTEM_API_DOC_PROMPT
                    doc_type_name = "api-doc"
                    print(f"[OpenAI-Doc] 文档类型: API接口文档 (api-doc)")
                elif body.document_type == "deployment-doc":
                    doc_system_prompt = SYSTEM_DEPLOYMENT_DOC_PROMPT
                    doc_type_name = "deployment-doc"
                    print(f"[OpenAI-Doc] 文档类型: 部署方案文档 (deployment-doc)")
                elif body.document_type == "overview-doc":
                    print(f"[OpenAI-Doc] 文档类型: 概要设计文档 (overview-doc)")
                else:
                    print(f"[OpenAI-Doc] 文档类型: 概要设计文档 (overview-doc)")
                    print(f"[文档类型-DEBUG-4] 使用默认类型, 因为传入值='{body.document_type}'不是已知的文档类型")
                
                print(f"[OpenAI-Doc] 使用文档类型: {doc_type_name}, 提示词长度: {len(doc_system_prompt)}")
                print(f"[文档类型-DEBUG-5] 最终确认使用的文档类型: '{doc_type_name}'")
                print(f"[文档类型-DEBUG-6] 使用的系统提示词类型: {doc_system_prompt[:50]}...")
                
                # 如果没有README，添加特殊指令
                if not has_readme:
                    first_system_prompt = first_system_prompt + "\n" + NO_README_INSTRUCTIONS
                
                if body.instructions:
                    # 将自定义指令添加到第一阶段提示
                    first_system_prompt = (
                        first_system_prompt
                        + "\n"
                        + ADDITIONAL_DOC_INSTRUCTIONS_PROMPT
                    )
                    # 将自定义指令添加到文档生成提示
                    doc_system_prompt = (
                        doc_system_prompt
                        + "\n"
                        + ADDITIONAL_DOC_INSTRUCTIONS_PROMPT
                    )

                # 第一阶段：获取项目架构解释（与图表生成相同）
                print("[OpenAI-Doc] 开始第一阶段: 生成项目架构解释")
                explanation = ""
                async for chunk in o3_service.call_o3_api_stream(
                    system_prompt=first_system_prompt,
                    data={
                        "file_tree": file_tree,
                        "readme": readme,
                        "instructions": body.instructions,
                    },
                    api_key=body.api_key,
                    reasoning_effort="medium",
                ):
                    explanation += chunk
                    # 流式返回解释块给客户端
                    yield f"data: {json.dumps({'status': 'explanation_chunk', 'chunk': chunk})}\n\n"
                
                print("[OpenAI-Doc] 第一阶段完成: 解释长度=" + str(len(explanation)))

                # 检查是否包含BAD_INSTRUCTIONS标记
                if "BAD_INSTRUCTIONS" in explanation:
                    yield f"data: {json.dumps({'error': 'Invalid or unclear instructions provided'})}\n\n"
                    return

                # 第二阶段：创建组件映射（与图表生成相同）
                print("[OpenAI-Doc] 开始第二阶段: 创建组件映射")
                full_second_response = ""
                async for chunk in o3_service.call_o3_api_stream(
                    system_prompt=SYSTEM_SECOND_PROMPT,
                    data={"explanation": explanation, "file_tree": file_tree},
                    api_key=body.api_key,
                    reasoning_effort="low",
                ):
                    full_second_response += chunk
                    # 流式返回映射块给客户端
                    yield f"data: {json.dumps({'status': 'mapping_chunk', 'chunk': chunk})}\n\n"
                
                print("[OpenAI-Doc] 第二阶段完成: 组件映射长度=" + str(len(full_second_response)))

                # 提取组件映射
                start_tag = "<component_mapping>"
                end_tag = "</component_mapping>"
                component_mapping_text = full_second_response[
                    full_second_response.find(start_tag) : full_second_response.find(
                        end_tag
                    ) + len(end_tag)
                ]

                # 第三阶段：生成详细文档
                print("[OpenAI-Doc] 开始第三阶段: 生成详细文档")
                document = ""
                async for chunk in o3_service.call_o3_api_stream(
                    system_prompt=doc_system_prompt,
                    data={
                        "explanation": explanation,
                        "component_mapping": component_mapping_text,
                        "instructions": body.instructions,
                    },
                    api_key=body.api_key,
                    reasoning_effort="high",  # 文档生成需要更高的推理努力
                ):
                    document += chunk
                    # 流式返回文档块给客户端
                    yield f"data: {json.dumps({'status': 'document_chunk', 'chunk': chunk})}\n\n"
                
                print("[OpenAI-Doc] 第三阶段完成: 文档长度=" + str(len(document)))

                # 检查是否包含BAD_INSTRUCTIONS标记
                if "BAD_INSTRUCTIONS" in document:
                    print("[OpenAI-Doc] 错误: 无效或不明确的指令")
                    yield f"data: {json.dumps({'error': 'Invalid or unclear instructions provided'})}\n\n"
                    return

                print(f"[OpenAI-Doc] 文档生成完成: 类型={doc_type_name}, 解释长度={len(explanation)}, 映射长度={len(component_mapping_text)}, 文档长度={len(document)}")

                # 发送最终结果给客户端
                yield f"data: {json.dumps({
                    'status': 'complete',
                    'document': document,
                    'explanation': explanation,
                    'mapping': component_mapping_text
                })}\n\n"

            except Exception as e:
                # 如果发生异常，返回错误信息
                print(f"[OpenAI-Doc] 生成过程中出错: {str(e)}")
                yield f"data: {json.dumps({'error': str(e)})}\n\n"

        # 返回流式响应对象
        return StreamingResponse(
            event_generator(),
            media_type="text/event-stream",
            headers={
                "X-Accel-Buffering": "no",
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
            },
        )
    except Exception as e:
        # 如果发生异常，返回错误信息
        print(f"[OpenAI-Doc] 请求处理出错: {str(e)}")
        return {"error": str(e)}


# 定义文档修改的路由
@router.post("/modify")
async def modify_document(request: Request, body: dict):
    """
    根据用户指令修改生成的文档
    
    Args:
        request (Request): FastAPI请求对象
        body (dict): 请求体，包含要修改的文档、指令和可选的API密钥
        
    Returns:
        dict: 包含修改后的文档或错误信息
    """
    print("[OpenAI-Doc] 文档修改请求")
    print(f"[OpenAI-Doc] 请求体字段: {list(body.keys())}")
    
    try:
        # 验证必要的请求字段
        required_fields = ["instructions", "current_document", "explanation"]
        missing_fields = [field for field in required_fields if field not in body]
        
        if missing_fields:
            print(f"[OpenAI-Doc] 缺少必要字段: {missing_fields}")
            raise HTTPException(status_code=400, detail=f"Missing required fields: {', '.join(missing_fields)}")
            
        # 获取请求数据
        instructions = body.get("instructions", "")
        current_document = body.get("current_document", "")
        explanation = body.get("explanation", "")
        api_key = body.get("api_key")
        
        print(f"[OpenAI-Doc] 修改指令长度: {len(instructions)}")
        print(f"[OpenAI-Doc] 当前文档长度: {len(current_document)}")
        print(f"[OpenAI-Doc] 解释文本长度: {len(explanation)}")
        print(f"[OpenAI-Doc] 是否提供API密钥: {api_key is not None}")
        
        # 检查关键数据是否为空
        if not instructions.strip():
            print("[OpenAI-Doc] 警告: 修改指令为空")
            return {"error": "Instructions cannot be empty"}
            
        if not current_document.strip() and not explanation.strip():
            print("[OpenAI-Doc] 错误: 当前文档和解释都为空")
            return {"error": "Both document and explanation cannot be empty"}
        
        # 设置系统提示词
        system_prompt = SYSTEM_DOC_MODIFY_PROMPT
        print(f"[OpenAI-Doc] 系统提示词长度: {len(system_prompt)}")
        
        # 准备API调用数据
        data = {
            "document": current_document,
            "instructions": instructions,
            "explanation": explanation
        }
        
        print("[OpenAI-Doc] 准备调用O3 API修改文档")
        
        # 调用OpenAI服务修改文档
        try:
            modified_document = await o3_service.call_o3_api(
                system_prompt=system_prompt,
                data=data,
                api_key=api_key,
                reasoning_effort="medium"
            )
            print(f"[OpenAI-Doc] O3 API调用成功，响应长度: {len(modified_document) if modified_document else 0}")
        except Exception as api_error:
            print(f"[OpenAI-Doc] O3 API调用失败: {str(api_error)}")
            return {"error": f"API service error: {str(api_error)}"}
        
        # 检查是否包含BAD_INSTRUCTIONS标记
        if "BAD_INSTRUCTIONS" in modified_document:
            print("[OpenAI-Doc] 接收到BAD_INSTRUCTIONS响应")
            return {"error": "Invalid or unclear instructions provided"}
            
        print("[OpenAI-Doc] 文档修改完成")
        
        # 返回修改后的文档
        return {"document": modified_document}
    
    except HTTPException as he:
        print(f"[OpenAI-Doc] HTTP异常: {he.detail}, 状态码: {he.status_code}")
        raise he
    except Exception as e:
        print(f"[OpenAI-Doc] 文档修改出错: {str(e)}")
        print(f"[OpenAI-Doc] 错误类型: {type(e).__name__}")
        import traceback
        print(f"[OpenAI-Doc] 错误堆栈: {traceback.format_exc()}")
        return {"error": str(e)} 